package com.ai.zuochengyun.phase01.class06;

import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Set;

public class Code07_Prim {
    public static Set<Edge> primMST(Graph graph) {
        // 解锁的边进入小根堆
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>();

        // 哪些点被解锁了
        Set<Node> nodeSet = new HashSet<>();
        Set<Edge> res = new HashSet<>();
        for (Node node : graph.nodes.values()) {
            if (!nodeSet.contains(node)) {
                nodeSet.add(node);
                // 解锁的边
                for (Edge edge : node.edges) {
                    priorityQueue.add(edge);
                }

                while (!priorityQueue.isEmpty()) {
                    Edge edge = priorityQueue.poll();
                    // 如果边对应的终点节点没有被解锁，那么就继续通过点找边
                    if (!nodeSet.contains(edge.to)) {
                        nodeSet.add(edge.to);
                        res.add(edge);

                        for (Edge nextEdge : edge.to.edges) {
                            priorityQueue.add(nextEdge);
                        }
                    }
                }
            }
        }
        return res;
    }
}
